#include "mpp.h"
#include "v4l2.h"
AVCodecContext *global_codec_ctx = nullptr; //H.264编码上下文(编解码)
AVFormatContext *global_fmt_ctx = nullptr;  //格式上下文 mp4容器
struct SwsContext *global_sws_ctx = nullptr; //格式转换上下文
AVFrame *global_frame = nullptr; //视频帧缓冲区
AVPacket *global_pkt = nullptr; //编码后的数据包
AVStream *global_stream = nullptr;

MppCtx mpp_ctx = NULL;
MppApi *mpi = NULL;
MppBufferGroup buf_grp = NULL;
MppBuffer frm_buf = NULL;
MppBuffer pkt_buf = NULL;

int rkmpp_init(const char* filename){
    // av_log_set_level(AV_LOG_DEBUG);

    MPP_RET ret = MPP_OK;
    MppPollType timeout = MPP_POLL_BLOCK;
    MppEncCfg cfg;
    size_t frame_size = WIDTH * HEIGHT * 3 / 2; //NV12 format

    //mpp context and api
    ret = mpp_create(&mpp_ctx, &mpi);
    if (MPP_OK != ret) {
        mpp_err("mpp_create failed\n");
        goto MPP_TEST_OUT;
    }

    //initial mpp encoder  MPP_CTX_ENC-->encoder\
                            MPP_VIDEO_CodingHEVC -->H265
    ret = mpp_init(mpp_ctx, MPP_CTX_ENC , MPP_VIDEO_CodingAVC);
    if(MPP_OK != ret){
        mpp_err("mpp_init failed\n");
        goto MPP_TEST_OUT;
    }

    // Initialize encoder configuration
    ret = mpp_enc_cfg_init(&cfg);
    if (MPP_OK != ret) {
        mpp_err("init encoder config failed\n");
        goto MPP_TEST_OUT;
    }

    // Get default configuration
    ret = mpi->control(mpp_ctx, MPP_ENC_GET_CFG, cfg);
    if (MPP_OK != ret) {
        mpp_err("get default encoder config failed\n");
        goto MPP_TEST_OUT;
    }    

    //配置编码器信息：
    //分辨率
    mpp_enc_cfg_set_s32(cfg, "prep:width", WIDTH);
    mpp_enc_cfg_set_s32(cfg, "prep:height", HEIGHT);
    mpp_enc_cfg_set_s32(cfg, "prep:hor_stride", WIDTH);
    mpp_enc_cfg_set_s32(cfg, "prep:ver_stride", HEIGHT);
    mpp_enc_cfg_set_s32(cfg, "prep:format", MPP_FMT_YUV420SP);

    // Rate control configuration
    // 码率、帧率
    mpp_enc_cfg_set_s32(cfg, "rc:mode", MPP_ENC_RC_MODE_CBR);
    mpp_enc_cfg_set_s32(cfg, "rc:bps_target", BITRATE);
    mpp_enc_cfg_set_s32(cfg, "rc:bps_max", BITRATE * 17 / 16);
    mpp_enc_cfg_set_s32(cfg, "rc:bps_min", BITRATE * 15 / 16);
    mpp_enc_cfg_set_s32(cfg, "rc:fps_in_num", FPS);
    mpp_enc_cfg_set_s32(cfg, "rc:fps_in_denom", 1);
    mpp_enc_cfg_set_s32(cfg, "rc:fps_out_num", FPS);
    mpp_enc_cfg_set_s32(cfg, "rc:fps_out_denom", 1);
    mpp_enc_cfg_set_s32(cfg, "rc:gop", FPS * 2); // 2 second GOP


    // Codec configuration
    // H.264
    mpp_enc_cfg_set_s32(cfg, "codec:type", MPP_VIDEO_CodingAVC);
    mpp_enc_cfg_set_s32(cfg, "h264:profile", 100); // High profile
    mpp_enc_cfg_set_s32(cfg, "h264:level", 40);    // Level 4.0
    mpp_enc_cfg_set_s32(cfg, "h264:cabac_en", 1);
    mpp_enc_cfg_set_s32(cfg, "h264:cabac_idc", 0);
    mpp_enc_cfg_set_s32(cfg, "h264:trans8x8", 1);

    // QP configuration
    mpp_enc_cfg_set_s32(cfg, "rc:qp_init", 26);
    mpp_enc_cfg_set_s32(cfg, "rc:qp_max", 51);
    mpp_enc_cfg_set_s32(cfg, "rc:qp_min", 10);
    mpp_enc_cfg_set_s32(cfg, "rc:qp_max_i", 51);
    mpp_enc_cfg_set_s32(cfg, "rc:qp_min_i", 10);
    mpp_enc_cfg_set_s32(cfg, "rc:qp_ip", 2);
    
    // set
    ret = mpi->control(mpp_ctx, MPP_ENC_SET_CFG, cfg);
    if(MPP_OK != ret){
        mpp_err("set encoder config failed\n");
        goto MPP_TEST_OUT;
    }

    //set output timeout
    ret = mpi->control(mpp_ctx, MPP_SET_OUTPUT_TIMEOUT, &timeout);
    if(MPP_OK != ret){
        mpp_err("set output timeout failed\n");
        goto MPP_TEST_OUT;
    }

    //create buffer group
    ret = mpp_buffer_group_get_internal(&buf_grp, MPP_BUFFER_TYPE_DRM);
    if(MPP_OK != ret){
        mpp_err("get buffer group failed\n");
        goto MPP_TEST_OUT;
    }

    //allocate frame buffer
    ret = mpp_buffer_get(buf_grp, &frm_buf, frame_size);
    if(MPP_OK != ret){
        mpp_err("get frame buffer failed\n");
        goto MPP_TEST_OUT;
    }

    //allocate packet buffer
    ret = mpp_buffer_get(buf_grp, &pkt_buf, frame_size);
    if(MPP_OK != ret){
        mpp_err("get packet buffer failed\n");
        goto MPP_TEST_OUT;
    }

    mpp_enc_cfg_deinit(cfg);

    // Initialize FFmpeg for MP4 muxing
    if (avformat_alloc_output_context2(&global_fmt_ctx, NULL, "mp4", filename) < 0) {
        std::cerr << "无法创建输出上下文" << std::endl;
        goto MPP_TEST_OUT;
    }

    global_stream = avformat_new_stream(global_fmt_ctx, NULL);
    if (!global_stream) {
        std::cerr << "无法创建视频流" << std::endl;
        goto FFMPEG_INIT_FAIL;
    }

    global_stream->codecpar->codec_id = AV_CODEC_ID_H264;
    global_stream->codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
    global_stream->codecpar->width = WIDTH;
    global_stream->codecpar->height = HEIGHT;
    global_stream->codecpar->format = AV_PIX_FMT_YUV420P;
    global_stream->codecpar->bit_rate = BITRATE;
    global_stream->time_base = (AVRational){1, FPS};

    if (avio_open(&global_fmt_ctx->pb, filename, AVIO_FLAG_WRITE) < 0) {
        std::cerr << "无法打开输出文件: " << filename << " - " << strerror(errno) << std::endl;
        goto FFMPEG_INIT_FAIL;
    }

    if (avformat_write_header(global_fmt_ctx, NULL) < 0) {
        std::cerr << "无法写入文件头" << std::endl;
        goto FFMPEG_INIT_FAIL;
    }

    // Initialize format conversion
    global_sws_ctx = sws_getContext(WIDTH, HEIGHT, AV_PIX_FMT_UYVY422,
                                    WIDTH, HEIGHT, AV_PIX_FMT_NV12,
                                    SWS_BILINEAR, NULL, NULL, NULL);
    if (!global_sws_ctx) {
        std::cerr << "无法初始化格式转换上下文" << std::endl;
        goto FFMPEG_INIT_FAIL;
    }

    return 0;

FFMPEG_INIT_FAIL:
    if (global_fmt_ctx) {
        avio_closep(&global_fmt_ctx->pb);
        avformat_free_context(global_fmt_ctx);
        global_fmt_ctx = nullptr;
    }

MPP_TEST_OUT:
    if (cfg) {
        mpp_enc_cfg_deinit(cfg);
    }
    if (frm_buf) {
        mpp_buffer_put(frm_buf);
    }
    if (pkt_buf) {
        mpp_buffer_put(pkt_buf);
    }
    if (buf_grp) {
        mpp_buffer_group_put(buf_grp);
    }
    if (mpp_ctx) {
        mpp_destroy(mpp_ctx);
    }
    return -1;

}

//编码、写入MP4
int encode_and_write_frame(uint8_t* frame_data, int frame_count) {
 
    AVFrame *av_frame = av_frame_alloc();
    if (!av_frame) {
        std::cerr << "无法分配帧" << std::endl;
        return -1;
    }
    av_frame->format = AV_PIX_FMT_NV12;
    av_frame->width = WIDTH;
    av_frame->height = HEIGHT;
    if (av_frame_get_buffer(av_frame, 32) < 0) {
        std::cerr << "无法分配帧缓冲区" << std::endl;
        av_frame_free(&av_frame);
        return -1;
    }

    uint8_t *src_data[1] = { frame_data };
    int src_linesize[1] = { WIDTH * 2 };
    if (!global_sws_ctx) {
        std::cerr << "格式转换上下文未初始化" << std::endl;
        av_frame_free(&av_frame);
        return -1;
    }
    sws_scale(global_sws_ctx, src_data, src_linesize, 0, HEIGHT, av_frame->data, av_frame->linesize);

    // Prepare MPP frame
    MppFrame mpp_frame = NULL;
    MPP_RET ret = mpp_frame_init(&mpp_frame);
    if (ret != MPP_OK) {
        mpp_err("Failed to init frame\n");
        av_frame_free(&av_frame);
        return -1;
    }

    void *frm_ptr = mpp_buffer_get_ptr_with_caller(frm_buf, __FUNCTION__);
    // NV12: Y plane followed by interleaved UV plane
    memcpy(frm_ptr, av_frame->data[0], WIDTH * HEIGHT); // Y plane
    memcpy(frm_ptr + WIDTH * HEIGHT, av_frame->data[1], WIDTH * HEIGHT / 2); // UV plane

    mpp_frame_set_width(mpp_frame, WIDTH);
    mpp_frame_set_height(mpp_frame, HEIGHT);
    mpp_frame_set_hor_stride(mpp_frame, WIDTH);
    mpp_frame_set_ver_stride(mpp_frame, HEIGHT);
    mpp_frame_set_fmt(mpp_frame, MPP_FMT_YUV420SP);
    mpp_frame_set_buffer(mpp_frame, frm_buf);
    mpp_frame_set_eos(mpp_frame, 0);

    // Encode frame
    MppPacket packet = NULL;
    ret = mpp_packet_init_with_buffer(&packet, pkt_buf);
    if (ret != MPP_OK) {
        mpp_err("Failed to init packet\n");
        mpp_frame_deinit(&mpp_frame);
        av_frame_free(&av_frame);
        return -1;
    }

    mpp_packet_set_length(packet, 0);
    MppMeta meta = mpp_frame_get_meta(mpp_frame);
    mpp_meta_set_packet(meta, KEY_OUTPUT_PACKET, packet);

    ret = mpi->encode_put_frame(mpp_ctx, mpp_frame);
    if (ret != MPP_OK) {
        mpp_err("Failed to put frame\n");
        mpp_packet_deinit(&packet);
        mpp_frame_deinit(&mpp_frame);
        av_frame_free(&av_frame);
        return -1;
    }

    mpp_frame_deinit(&mpp_frame);
    av_frame_free(&av_frame);

    ret = mpi->encode_get_packet(mpp_ctx, &packet);
    if (ret != MPP_OK) {
        mpp_err("Failed to get packet\n");
        mpp_packet_deinit(&packet);
        return -1;
    }

    if (mpp_packet_get_length(packet) > 0) {
        AVPacket *av_pkt = av_packet_alloc();
        if (!av_pkt) {
            mpp_err("Failed to allocate AVPacket\n");
            mpp_packet_deinit(&packet);
            return -1;
        }

        av_pkt->data = (uint8_t*)mpp_packet_get_data(packet);
        av_pkt->size = mpp_packet_get_length(packet);
        av_pkt->stream_index = global_stream->index;
        av_pkt->pts = av_rescale_q(frame_count, (AVRational){1, FPS}, global_stream->time_base);
        av_pkt->dts = av_pkt->pts;

        if (av_interleaved_write_frame(global_fmt_ctx, av_pkt) < 0) {
            std::cerr << "写入数据包失败" << std::endl;
            av_packet_free(&av_pkt);
            mpp_packet_deinit(&packet);
            return -1;
        }

        av_packet_free(&av_pkt);
    }

    mpp_packet_deinit(&packet);
    return 0;
}

//拍照
int save_snapshot(uint8_t * frame_data, const char* filename){

    // Initialize JPEG encoder 
    const AVCodec *jpeg_codec = avcodec_find_encoder(AV_CODEC_ID_MJPEG);
    if (!jpeg_codec) {
        std::cerr << "未找到 JPEG 编码器" << std::endl;
        return -1;
    }

    // 编码器上下文
    AVCodecContext *jpeg_ctx = avcodec_alloc_context3(jpeg_codec);
    if(!jpeg_ctx){
        std::cerr << "无法分配JPEG编码上下文" << std::endl;
        return -1;
    }
    jpeg_ctx->width = WIDTH;
    jpeg_ctx->height = HEIGHT;
    jpeg_ctx->pix_fmt = AV_PIX_FMT_YUVJ420P;
    jpeg_ctx->time_base = (AVRational){1, FPS};
    av_opt_set_int(jpeg_ctx->priv_data, "q", 2, 0);

    //打开编码器
    if (avcodec_open2(jpeg_ctx, jpeg_codec, NULL) < 0) {
        std::cerr << "无法打开JPEG编码器" << std::endl;
        avcodec_free_context(&jpeg_ctx);
        return -1;
    }

    // Convert UYVY to YUV420P
    SwsContext *sws_ctx = sws_getContext(WIDTH, HEIGHT, AV_PIX_FMT_UYVY422,
                        WIDTH, HEIGHT, AV_PIX_FMT_YUV420P,
                        SWS_BILINEAR, NULL ,NULL ,NULL);
    
    if (!sws_ctx) {
        std::cerr << "无法初始化拍照转换上下文" << std::endl;
        avcodec_free_context(&jpeg_ctx);
        return -1;
    }

    // --- 帧分配 ---
    AVFrame *frame = av_frame_alloc();
    if (!frame) {
        std::cerr << "无法分配帧" << std::endl;
        sws_freeContext(sws_ctx);
        avcodec_free_context(&jpeg_ctx);
        return -1;
    }
    frame->format = AV_PIX_FMT_YUV420P;
    frame->width = WIDTH;
    frame->height = HEIGHT;
    if (av_frame_get_buffer(frame, 32) < 0) {  //设置好宽度、高度、格式，自动为AVFrame分配空间 32字节对齐
        std::cerr << "无法分配帧缓冲区" << std::endl;
        av_frame_free(&frame);
        sws_freeContext(sws_ctx);
        avcodec_free_context(&jpeg_ctx);
        return -1;
    }

    uint8_t *src_data[1] = {frame_data};
    int src_linesize[1] = {WIDTH * 2};
    //像素格式分辨率转换     in           in                    out             out
    sws_scale(sws_ctx, src_data, src_linesize ,0 ,HEIGHT, frame->data, frame->linesize);

    //Encode frame
    AVPacket *pkt = av_packet_alloc();
    if (!pkt) {
        std::cerr << "无法分配数据包" << std::endl;
        av_frame_free(&frame);
        sws_freeContext(sws_ctx);
        avcodec_free_context(&jpeg_ctx);
        return -1;
    }

    // 发送帧(frame)到编码器
    if (avcodec_send_frame(jpeg_ctx, frame) < 0) {
        std::cerr << "发送帧到编码器失败" << std::endl;
        av_packet_free(&pkt);
        av_frame_free(&frame);
        sws_freeContext(sws_ctx);
        avcodec_free_context(&jpeg_ctx);
        return -1;
    }

    // 接收编码后的数据包
    if(avcodec_receive_packet(jpeg_ctx, pkt) < 0){
        std::cerr << "Failed to receive JPEG packet" << std::endl;
        av_packet_free(&pkt);
        av_frame_free(&frame);
        sws_freeContext(sws_ctx);
        avcodec_free_context(&jpeg_ctx);
        return -1;
    }

    //write to file
    FILE *file = fopen(filename, "wb");
    if(!file){
        std::cerr << "Failed to open snapshot file: " << strerror(errno) << std::endl;
        av_packet_free(&pkt);
        av_frame_free(&frame);
        sws_freeContext(sws_ctx);
        avcodec_free_context(&jpeg_ctx);
        return -1;
    }
    std::cout<< "open snapshot file ok" <<std::endl;
    //压缩编码的数据  被写入的每个元素的大小，以字节为单位   数据的大小
    fwrite(pkt->data, 1, pkt->size, file);
    fclose(file);

    // Clean up
    av_packet_free(&pkt);
    av_frame_free(&frame);
    sws_freeContext(sws_ctx);
    avcodec_free_context(&jpeg_ctx);
    return 0;
    

}

void mpp_cleanup(void) {
    // Flush MPP encoder
    if (mpi && mpp_ctx) {
        MppFrame frame = NULL;
        MPP_RET ret = mpp_frame_init(&frame);
        if (ret == MPP_OK) {
            mpp_frame_set_eos(frame, 1);
            mpi->encode_put_frame(mpp_ctx, frame);
            mpp_frame_deinit(&frame);

            while (1) {
                MppPacket packet = NULL;
                ret = mpp_packet_init_with_buffer(&packet, pkt_buf);
                if (ret != MPP_OK)
                    break;

                ret = mpi->encode_get_packet(mpp_ctx, &packet);
                if (ret != MPP_OK) {
                    mpp_packet_deinit(&packet);
                    break;
                }

                if (mpp_packet_get_length(packet) > 0 && global_fmt_ctx) {
                    AVPacket *av_pkt = av_packet_alloc();
                    if (av_pkt) {
                        av_pkt->data = (uint8_t*)mpp_packet_get_data(packet);
                        av_pkt->size = mpp_packet_get_length(packet);
                        av_pkt->stream_index = global_stream->index;
                        av_pkt->pts = av_pkt->dts = AV_NOPTS_VALUE;

                        av_interleaved_write_frame(global_fmt_ctx, av_pkt);
                        av_packet_free(&av_pkt);
                    }
                }

                if (mpp_packet_get_eos(packet)) {
                    mpp_packet_deinit(&packet);
                    break;
                }

                mpp_packet_deinit(&packet);
            }
        }
    }

    // Write MP4 trailer
    if (global_fmt_ctx) {
        av_write_trailer(global_fmt_ctx);
    }

     // Cleanup MPP resources
     if (mpi) {
        mpi->reset(mpp_ctx);
    }
    
    if (mpp_ctx) {
        mpp_destroy(mpp_ctx);
        mpp_ctx = NULL;
    }
    
    if (frm_buf) {
        mpp_buffer_put(frm_buf);
        frm_buf = NULL;
    }
    
    if (pkt_buf) {
        mpp_buffer_put(pkt_buf);
        pkt_buf = NULL;
    }
    
    if (buf_grp) {
        mpp_buffer_group_put(buf_grp);
        buf_grp = NULL;
    }

    // Clean up FFmpeg resources
    if (global_sws_ctx) {
        sws_freeContext(global_sws_ctx);
        global_sws_ctx = nullptr;
    }
    if (global_fmt_ctx) {
        avio_closep(&global_fmt_ctx->pb);
        avformat_free_context(global_fmt_ctx);
        global_fmt_ctx = nullptr;
    }

    mpp_log("Cleanup completed\n");
}

